import os
import dspy
from dotenv import load_dotenv

# Default deployment to use if not specified in environment variables
DEFAULT_DEPLOYMENT = "gpt-4o-mini"

# Load environment variables
load_dotenv(".env.local")
# Fallback to .env if .env.local doesn't exist
if not os.path.exists(".env.local") and os.path.exists(".env"):
    load_dotenv(".env")


# Create DSPy language model
# Instead of passing the client directly, provide the config parameters
azure_llm = dspy.LM(f'azure/{DEFAULT_DEPLOYMENT}')
dspy.settings.configure(lm=azure_llm)

# Define a simple question-answering module using DSPy
class QAModule(dspy.Module):
    def __init__(self):
        super().__init__()
        self.generate_answer = dspy.ChainOfThought("question -> answer")
    
    def forward(self, question):
        return self.generate_answer(question=question)

# Example usage
def main():
    # Check if environment variables are set
    
    # Create the QA module
    qa_module = QAModule()
    
    # Example questions
    questions = [
        "What are the main benefits of using DSPy?",
        "How does DSPy integrate with Azure OpenAI?",
        "What are some common use cases for DSPy?"
    ]
    
    print("\n" + "="*50)
    print("DSPy + Azure OpenAI Example")
    print("="*50)
    print(f"Using deployment: {DEFAULT_DEPLOYMENT}")
    print("="*50 + "\n")
    
    # Generate answers for each question
    for question in questions:
        print(f"\nQuestion: {question}")
        try:
            response = qa_module(question)
            print(f"Answer: {response.answer}")
        except Exception as e:
            print(f"Error: {str(e)}")
            print("Check your credentials and connection to Azure OpenAI.")
            return

if __name__ == "__main__":
    main() 